import os
import os.path as osp
from easydict import EasyDict as edict
import socket
import logging

def add_pathdir(config):
    hostname    = socket.gethostname()
    if 'tzz-PC' in hostname:
        person_name     = 'tzz'
    elif 'zws-625' in hostname or 'wgw-1018' in hostname:
        person_name     = 'hzw'
    else:
        assert False, "hostname:{} not seen before".format(hostname)
    cur_file_path       = osp.abspath(__file__)
    project_dir         = osp.dirname(osp.dirname(osp.dirname(cur_file_path)))
    print("project_dir is:{}".format(project_dir))
    logging.info("project_dir is:{}".format(project_dir))
    if 'pathdir' not in config:
        config.pathdir  = edict()

    config.person_name  = person_name
    config.pathdir.local_home_dir   = cur_file_path.split(config.person_name)[0] + config.person_name
    config.pathdir.project_dir      = project_dir
    return config

def add_dataset_params(config, train_dataset_list=None, test_dataset_list=None,
    dataset_type='det', do_copy_image=False):
    """
    modify from smart
    """
    if 'dataset' not in config:
        config.dataset  = edict()
    # train
    config.dataset.train_roidb_path_list    = []
    config.dataset.train_imglst_path_list   = []
    config.dataset.train_imgidx_path_list   = []
    config.dataset.train_imgrec_path_list   = []
    config.dataset.train_coco_annotation_path_list  = []
    if train_dataset_list is not None:
        for dataset in train_dataset_list:
            # train_dataset_list is ['coco#train2017']
            dataset_name, imageset_name     = dataset.split("#")
            roidb_path, imglst_path, imgidx_path, imgrec_path, anno_path   = \
                _get_roidb_image_path(config, dataset_name, imageset_name, dataset_type, do_copy_image=do_copy_image)
            config.dataset.train_roidb_path_list.append(roidb_path)
            config.dataset.train_imglst_path_list.append(imglst_path)
            config.dataset.train_imgidx_path_list.append(imgidx_path)
            config.dataset.train_imgrec_path_list.append(imgrec_path)
            config.dataset.train_coco_annotation_path_list.append(anno_path)

    # test
    config.dataset.test_roidb_path_list     = []
    config.dataset.test_imglst_path_list    = []
    config.dataset.test_imgidx_path_list    = []
    config.dataset.test_imgrec_path_list    = []
    config.dataset.test_coco_annotation_path_list  = []
    if test_dataset_list is not None:
        for dataset in test_dataset_list:
            # train_dataset_list is ['coco#val2017']
            dataset_name, imageset_name     = dataset.split("#")
            roidb_path, imglst_path, imgidx_path, imgrec_path, anno_path   = \
                _get_roidb_image_path(config, dataset_name, imageset_name, dataset_type, do_copy_image=do_copy_image)
            config.dataset.test_roidb_path_list.append(roidb_path)
            config.dataset.test_imglst_path_list.append(imglst_path)
            config.dataset.test_imgidx_path_list.append(imgidx_path)
            config.dataset.test_imgrec_path_list.append(imgrec_path)
            config.dataset.test_coco_annotation_path_list.append(anno_path)

    return config

def _get_roidb_image_path(config, dataset_name, imageset_name, dataset_type, do_copy_image=False):
    """
    helper function for add_dataset_params
    """
    if dataset_name     == 'coco':
        dataset_path    = 'projects/dataset/coco/'
    elif dataset_name   == 'voc':
        dataset_name    = 'projects/dataset/VOCdevkit'
    elif dataset_name   == 'kitti':
        # TODO
        pass
    elif dataset_name   == 'dota':
        # TODO
        pass
    else:
        raise ValueError("Unknown dataset name {}".format(dataset_name))
    
    if dataset_name     == 'dota':
        # using roidb data
        setName, labelName  = imageset_name.split("_")
        roidb_path      = osp.join(config.pathdir.local_home_dir, dataset_path, \
             'roidbs', '%s_%s_gt_roidb_%s.pkl' % (setName, dataset_type, labelName))
    else:
        roidb_path      = osp.join(config.pathdir.local_home_dir, dataset_path, 'roidbs', '%s_%s_gt_roidb.pkl' % (imageset_name, dataset_type))
    anno_dir            = osp.join(config.pathdir.local_home_dir, dataset_path, 'annotations')
    assert dataset_type == 'det', "dataset_type:{} must be det now".format(dataset_type)
    anno_path           = osp.join(anno_dir, 'instances_{}.json'.format(imageset_name))
    assert do_copy_image is False, "not considering copy image here"
    imglstrec_dir       = osp.join(config.pathdir.local_home_dir, dataset_path, 'images_lst_rec')
    imglst_path         = osp.join(imglstrec_dir, '{}.lst'.format(imageset_name))
    imgidx_path         = osp.join(imglstrec_dir, '{}.idx'.format(imageset_name))
    imgrec_path         = osp.join(imglstrec_dir, '{}.rec'.format(imageset_name))
    return roidb_path, imglst_path, imgidx_path, imgrec_path, anno_path
